package com.graduation.bird.taskmanager;

import org.springframework.stereotype.Service;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import java.io.IOException;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.CompletableFuture;

@Service
public class TaskManager {
    private final ConcurrentHashMap<String, CompletableFuture<String>> tasks = new ConcurrentHashMap<>();
    private final ConcurrentHashMap<String, SseEmitter> emitters = new ConcurrentHashMap<>();

    public String startTask(CompletableFuture<String> future) {
        String taskId = generateTaskId();
        tasks.put(taskId, future);

        // 监听任务完成并推送 SSE
        future.thenAccept(result -> {
            SseEmitter emitter = emitters.get(taskId);
            if (emitter != null) {
                try {
                    emitter.send(SseEmitter.event()
                            .name("complete")
                            .data(result));
                    emitter.complete();
                } catch (IOException e) {
                    emitter.completeWithError(e);
                } finally {
                    emitters.remove(taskId);
                }
            }
        });

        return taskId;
    }

    public void registerEmitter(String taskId, SseEmitter emitter) {
        emitters.put(taskId, emitter);
    }

    public String getTaskResult(String taskId) {
        CompletableFuture<String> future = tasks.get(taskId);
        if (future != null && future.isDone()) {
            try {
                return future.get();
            } catch (Exception e) {
                return "failed";
            }
        }
        return null;
    }

    private String generateTaskId() {
        return Long.toHexString(System.currentTimeMillis());
    }
}

